package sdp.server;

import static com.googlecode.javacv.cpp.opencv_core.*;
import static com.googlecode.javacv.cpp.opencv_video.*;

import sdp.server.math.MathUtils;
import sdp.server.math.Vector2;

/** Stores the state of a robot - either the player or the opponent:
 * 	- The robot's team;
 *  - The robot's position and rotation;
 *  - The robot's size.
 *  
 *  @author Cristian Cobzarenco
 */
public class RobotState {
	private Team     team;
	private Vector2  position, velocity;
	private double   rotation, angular;
	private Vector2  size;
	private CvKalman kalman;
	private CvMat    measured;
	private CvMat    prediction;
	private CvMat    transitionMatrix;
	private boolean  useKalman = true;
	private int      undetectedFrameCount = 0;
		
	/** Enum representing the two possible 'teams' in game: yellow or blue. */
	public enum Team {
		YELLOW, BLUE;
		
		public Team getOtherTeam() {
			if( this == YELLOW )
				return BLUE;
			else
				return YELLOW;
		}
	}
	
	/** Construct a robot state, specifying team and size. The position and rotation are set to NaN, waiting
	 * the first update. The velocities are set to zero.
	 * 
	 * @param team The team the robot is on.
	 * @param size The size of the robot [width, height] in metres.
	 */
	public RobotState( Team team, Vector2 size ) {
		this( team, size, Vector2.NaN, Double.NaN );
	}
	
	/** Construct a robot state, specifying team and size as well as the initial position and rotation. The
	 * velocities are set to zero.
	 * 
	 * @param team The team the robot is on.
	 * @param size The size of the robot [width, height] in metres.
	 * @param position The position of the robot [ x, y ] in metres.
	 * @param rotation The angle the robot's direction makes with the horizontal axis.
	 */
	public RobotState( Team team, Vector2 size, Vector2 position, double rotation ) {
		assert size != null && position != null;
		kalman = cvCreateKalman( 6, 3, 0 );
		measured   = CvMat.create( 3, 1, CV_32FC1 );
		prediction = CvMat.create( 6, 1, CV_32FC1 );
		
		setProcessNoise( 1e-3 );
		setMeasurementNoise( 1e-3 );
		setPostError( 1e-1 );
		
		setupTransitionMatrix();
		
		this.team = team;
		this.size = size;
		reset( position, rotation );
	}
	
	/** Set the process noise covariance used by the Kalman filter.
	 * 
	 * @param newProcessNoise The new value for the process noise.
	 */
	public void setProcessNoise( double newProcessNoise ) {
		cvSetIdentity( kalman.process_noise_cov(), cvScalarAll( newProcessNoise ) );
	}
	
	/** Set the measurement noise used by the Kalman filter.
	 * 
	 * @param newMeasurementNoise The new value for the measurement noise.
	 */
	public void setMeasurementNoise( double newMeasurementNoise ) {
		cvSetIdentity( kalman.measurement_noise_cov(), cvScalarAll( newMeasurementNoise ) );
	}
	
	
	/** Set the post error covariance used by the Kalman filter.
	 * 
	 * @param newPostError The new post error covariance value.
	 */
	public void setPostError( double newPostError ) {
		cvSetIdentity( kalman.error_cov_post(), cvScalarAll( newPostError ) );
	}
	
	/** Reset the robot state. This should be called when a misdetection has occurred in the previous time steps and
	 * we cannot the robot's velocity, but we know its currently detected position.
	 * 
	 * The velocities will be set to zero.
	 * 
	 * @param newPosition The new detected position the robot in metres.
	 * @param newRotation The new detected rotation of the robot in radians.
	 */
	public void reset( Vector2 newPosition, double newRotation ) {	
		angular   = 0.0;
		velocity  = Vector2.ZERO;
		position  = newPosition;
		rotation  = MathUtils.capAngle( newRotation );
		
		if( !newPosition.isNaN() && !Double.isNaN( newRotation ) ) {
			kalman.state_pre().put( 0, position.getX() );
			kalman.state_pre().put( 1, position.getY() );
			kalman.state_pre().put( 2, rotation );
			kalman.state_pre().put( 3, 0.0 );
			kalman.state_pre().put( 4, 0.0 );
			kalman.state_pre().put( 5, 0.0 );
			
			cvSetIdentity( kalman.measurement_matrix() );
		}		
	}
	
	/** Returns the robot's linear velocity. This is automatically computed by update().
	 * 
	 * @return The robot's linear velocity in m/s.
	 */
	public Vector2 getLinearVelocity() {
		return velocity;
	}
	
	/** Returns the robot's angular velocity. This is automatically computed by update().
	 * 
	 * @return The robot's angular velocity in rad/s.
	 */
	public double getAngularVelocity() {
		return angular;
	}
	
	/** Returns the team the robot is on.
	 * 
	 * @return The team the robot is on.
	 */
	public Team getTeam() {
		return team;
	}

	/** Returns the size of the robot. This is set at construction.
	 * 
	 * @return The size of the robot [width, height] in metres.
	 */
	public Vector2 getSize() {
		return size;
	}
	
	/** Returns the position of the robot. This is set on update().
	 * 
	 * @return The position of the robot [x, y] in metres. NaN if the position is not known yet.
	 */
	public Vector2 getPosition() {
		return position;
	}
	
	/** Returns the angle the robot's 'T' makes with the horizontal axis.
	 * 
	 * @return The robot's orientation in radians. NaN if the orientation is not known yet.
	 */
	public double getRotation() {
		return MathUtils.capAngle( rotation );
	}
	
	/** Returns true if the robot cannot be detected. The state is still predicted
	 * if true.
	 * 
	 * @return True if the robot could not be detected, false otherwise.
	 */
	public boolean isDetected() {
		return undetectedFrameCount <= 5 && !position.isNaN();
	}
	
	/** Update the robot's position and rotation. The velocity and angular velocity are also computed using Euler
	 * integration: (newPosition - oldPosition) / timeStep. 
	 * 
	 * @param newPosition The new detected position.
	 * @param newRotation The new detected rotation.
	 * @param timeStep The time elapsed since the previous update.
	 */
	public void update( Vector2 newPosition, double newRotation, double timeStep ) {
		assert newPosition == null || !newPosition.isNaN();
		assert !Double.isNaN( newRotation );
		assert timeStep > 0.0;
		
		if( position.isNaN() || Double.isNaN( rotation ) ) {
			if( newPosition != null )
				reset( newPosition, newRotation );
		} else if( newPosition != null ) {
			undetectedFrameCount = 0;
			if( useKalman ) {
				cvKalmanPredict( kalman, prediction );
			
				newRotation = rotation + MathUtils.angleDiff( MathUtils.capAngle( newRotation ), MathUtils.capAngle( rotation ) );
				assert Math.abs( rotation - newRotation ) <= Math.PI * 2.0 : "rotation = " + rotation + ", newRotation = " + newRotation;
				
				measured.put( 0, newPosition.getX() );
				measured.put( 1, newPosition.getY() );
				measured.put( 2, newRotation );
				
				CvMat estimated  = cvKalmanCorrect( kalman, measured );
				position = new Vector2( estimated.get( 0 ), estimated.get( 1 ) );
				rotation = estimated.get( 2 );
				velocity = new Vector2( prediction.get( 3 ) / timeStep, prediction.get( 4 ) / timeStep );
				angular  = prediction.get( 5 ) / timeStep;
				
				if( angular > Math.PI * 2.0 ) {
					kalman.state_pre().put(5, 0.0);
					kalman.state_pre().put(3, newRotation);
					angular = 0.0;
				}
				
				if( velocity.getX() > 10.0 || velocity.getY() > 10.0 || velocity.getX() < -10.0 || velocity.getY() < -10.0 )
					velocity = Vector2.ZERO;
			} else {
				position = newPosition;
				rotation = newRotation;
				
				velocity = newPosition.minus( position ).times( 1.0 / timeStep );
				angular  = MathUtils.angleDiff( newRotation, rotation ) / timeStep;
			}
		} else {
			++ undetectedFrameCount;
			position = position.plus( velocity.times( timeStep ) );
			velocity = velocity.times( 0.9 );
			angular *= 0.9;
		}
	}
	
	/** Return a clone of this RobotState. */
	@Override
	public RobotState clone() {
		RobotState r = new RobotState( team, size, position, rotation );
		r.angular  = angular;
		r.velocity = velocity;
		return r;
	}
	
	private void setupTransitionMatrix() {
		transitionMatrix = CvMat.create( 6, 6, CV_32FC1 );
		transitionMatrix.put( 0, 0, 1 );
		transitionMatrix.put( 0, 1, 0 );
		transitionMatrix.put( 0, 2, 0 );
		transitionMatrix.put( 0, 3, 1 );
		transitionMatrix.put( 0, 4, 0 );
		transitionMatrix.put( 0, 5, 0 );
		
		transitionMatrix.put( 1, 0, 0 );
		transitionMatrix.put( 1, 1, 1 );
		transitionMatrix.put( 1, 2, 0 );
		transitionMatrix.put( 1, 3, 0 );
		transitionMatrix.put( 1, 4, 1 );
		transitionMatrix.put( 1, 5, 0 );
		
		transitionMatrix.put( 2, 0, 0 );
		transitionMatrix.put( 2, 1, 0 );
		transitionMatrix.put( 2, 2, 1 );
		transitionMatrix.put( 2, 3, 0 );
		transitionMatrix.put( 2, 4, 0 );
		transitionMatrix.put( 2, 5, 1 );
		
		transitionMatrix.put( 3, 0, 0 );
		transitionMatrix.put( 3, 1, 0 );
		transitionMatrix.put( 3, 2, 0 );
		transitionMatrix.put( 3, 3, 1 );
		transitionMatrix.put( 3, 4, 0 );
		transitionMatrix.put( 3, 5, 0 );
		
		transitionMatrix.put( 4, 0, 0 );
		transitionMatrix.put( 4, 1, 0 );
		transitionMatrix.put( 4, 2, 0 );
		transitionMatrix.put( 4, 3, 0 );
		transitionMatrix.put( 4, 4, 1 );
		transitionMatrix.put( 4, 5, 0 );
		
		transitionMatrix.put( 5, 0, 0 );
		transitionMatrix.put( 5, 1, 0 );
		transitionMatrix.put( 5, 2, 0 );
		transitionMatrix.put( 5, 3, 0 );
		transitionMatrix.put( 5, 4, 0 );
		transitionMatrix.put( 5, 5, 1 );
		
		kalman.transition_matrix( transitionMatrix );
	}
}
